import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from deap import base, creator, tools, algorithms
import random

r_I = 0.1
K_I = 1000
r_P = 0.05
K_P = 1000
r_B = 0.02
K_B = 1000
alpha_I = 0.005
alpha_P = 0.01
gamma_F = 0.02
delta_B = 0.1
R_P = 5
C_P = 2
C_F = 1

def ecosystem_model(y, t, F):
    I, P, B = y
    dI_dt = r_I * I * (1 - I / K_I) - alpha_I * P - 0.1 * F
    dP_dt = r_P * P * (1 - P / K_P) + gamma_F * F - alpha_P * I
    dB_dt = r_B * B * (1 - B / K_B) + delta_B * P
    return [dI_dt, dP_dt, dB_dt]

def profit(F, P):
    return R_P * P - C_P * P - C_F * F

def pest_control(I, P):
    return I / P

def biodiversity(B):
    return B

creator.create("FitnessMulti", base.Fitness, weights=(1.0, -1.0, - 1.0))
creator.create("Individual", list, fitness=creator.FitnessMulti)

def create_individual():
    return [random.uniform(0, 1)]

toolbox = base.Toolbox()
toolbox.register("individual", tools.initIterate, creator.Individual, create_individual)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

def evaluate(individual):
    F = individual[0]
    y0 = [50, 500, 200]
    t = np.linspace(0, 100, 1000)
    sol = odeint(ecosystem_model, y0, t, args=(F,))
    I, P, B = sol.T
    profit_value = profit(F, P[-1])
    pest_control_value = pest_control(I[-1], P[-1])
    biodiversity_value = biodiversity(B[-1])
    return profit_value, pest_control_value, biodiversity_value

toolbox.register("mate", tools.cxBlend, alpha=0.5)
toolbox.register("mutate", tools.mutGaussian, mu=0.5, sigma=0.1, indpb=0.2)
toolbox.register("select", tools.selNSGA2)
toolbox.register("evaluate", evaluate)

def run_ga():
    population = toolbox.population(n=50)
    generations = 300
    for gen in range(generations):
        offspring = list(map(toolbox.clone, population))
        for child1, child2 in zip(offspring[::2], offspring[1::2]):
            if random.random() < 0.7:
                toolbox.mate(child1, child2)
                del child1.fitness.values
                del child2.fitness.values
        for mutant in offspring:
            if random.random() < 0.2:
                toolbox.mutate(mutant)
                del mutant.fitness.values
        invalid_individuals = [ind for ind in offspring if not ind.fitness.valid]
        fitnesses = list(map(toolbox.evaluate, invalid_individuals))
        for ind, fit in zip(invalid_individuals, fitnesses):
            ind.fitness.values = fit
        population[:] = toolbox.select(offspring, len(population))
        fits = [ind.fitness.values[0] for ind in population]
        print(f"Generation {gen}: Max profit: {max(fits)}")
    return population


final_population = run_ga()

def plot_results(population):
    profits = [ind.fitness.values[0] for ind in population]
    pest_controls = [ind.fitness.values[1] for ind in population]
    biodiversities = [ind.fitness.values[2] for ind in population]
    plt.figure(figsize=(12, 6))
    plt.scatter(profits, pest_controls, c=biodiversities, cmap='viridis')
    plt.title("Profit vs Pest Control (Colored by Biodiversity)")
    plt.xlabel("Profit")
    plt.ylabel("Pest Control")
    plt.colorbar(label="Biodiversity")
    plt.show()

plot_results(final_population)
